from fastapi import APIRouter, Depends, HTTPException
from typing import List, Dict, Optional, Any
from pydantic import BaseModel

# 创建API路由实例
router = APIRouter()

# 示例：模型加载请求模型
class ModelLoadRequest(BaseModel):
    model_name: str
    device: Optional[str] = None
    use_quantization: Optional[bool] = False

# 示例：文本生成请求模型
class TextGenerationRequest(BaseModel):
    model_name: Optional[str] = None
    prompt: str
    max_tokens: Optional[int] = 2048
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.95
    stream: Optional[bool] = False

# 示例：微调请求模型
class FineTuneRequest(BaseModel):
    model_name: str
    dataset_path: str
    output_dir: Optional[str] = None
    num_epochs: Optional[int] = 3
    batch_size: Optional[int] = 4
    learning_rate: Optional[float] = 2e-5

# 示例：RAG查询请求模型
class RAGQueryRequest(BaseModel):
    query: str
    model_name: Optional[str] = None
    top_k: Optional[int] = 5
    use_rag: Optional[bool] = True

# 健康检查接口
@router.get("/health", summary="服务健康检查")
def health_check():
    """检查服务是否正常运行"""
    return {
        "status": "healthy",
        "message": "服务正常运行中"
    }

# 列出可用模型接口
@router.get("/models", summary="列出可用模型")
def list_models():
    """获取当前已加载的所有模型"""
    from models.model_manager import ModelManager
    from config import settings
    
    model_manager = ModelManager()
    available_models = model_manager.get_available_models()
    
    # 添加额外信息，帮助前端理解模型列表状态
    extra_info = {
        "has_models": len(available_models) > 0,
        "ollama_enabled": settings.ENABLE_OLLAMA,
    }
    
    # 如果启用了Ollama但没有模型，添加提示信息
    if settings.ENABLE_OLLAMA and len(available_models) == 0:
        extra_info["hint"] = "Ollama服务已启用，但未检测到可用模型。请确保Ollama服务运行正常且已下载模型。"
    
    return {
        "models": available_models,
        "default_model": model_manager.get_default_model_name(),
        "extra_info": extra_info
    }

# 加载模型接口
@router.post("/models/load", summary="加载模型")
def load_model(request: ModelLoadRequest):
    """加载指定的模型"""
    from models.model_manager import ModelManager
    try:
        result = ModelManager().load_model(
            model_name=request.model_name,
            device=request.device,
            use_quantization=request.use_quantization
        )
        return {
            "status": "success",
            "message": f"模型 {request.model_name} 加载成功",
            "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型加载失败: {str(e)}")

# 文本生成接口
@router.post("/generate", summary="文本生成")
def generate_text(request: TextGenerationRequest):
    """使用模型生成文本"""
    from models.model_manager import ModelManager
    try:
        result = ModelManager().generate_text(
            prompt=request.prompt,
            model_name=request.model_name,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            stream=request.stream
        )
        return {
            "status": "success",
            "message": "文本生成成功",
            "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"文本生成失败: {str(e)}")

# 微调模型接口
@router.post("/fine-tune", summary="微调模型")
def fine_tune_model(request: FineTuneRequest):
    """微调指定的模型"""
    from fine_tune.fine_tuner import FineTuner
    try:
        tuner = FineTuner()
        result = tuner.fine_tune(
            model_name=request.model_name,
            dataset_path=request.dataset_path,
            output_dir=request.output_dir,
            num_epochs=request.num_epochs,
            batch_size=request.batch_size,
            learning_rate=request.learning_rate
        )
        return {
            "status": "success",
            "message": f"模型 {request.model_name} 微调成功",
            "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型微调失败: {str(e)}")

# RAG查询接口
@router.post("/rag/query", summary="RAG检索查询")
def rag_query(request: RAGQueryRequest):
    """使用RAG检索增强进行查询"""
    from rag.rag_manager import RAGManager
    try:
        manager = RAGManager()
        result = manager.query(
            query=request.query,
            model_name=request.model_name,
            top_k=request.top_k,
            use_rag=request.use_rag
        )
        return {
            "status": "success",
            "message": "RAG查询成功",
            "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"RAG查询失败: {str(e)}")

# 上传知识库文档接口
@router.post("/rag/upload", summary="上传知识库文档")
def upload_knowledge_file(file: Optional[Any] = None):
    """上传文档到知识库"""
    from rag.rag_manager import RAGManager
    try:
        # 这里简化处理，实际应用中需要处理文件上传
        manager = RAGManager()
        # result = manager.add_document(file)
        return {
            "status": "success",
            "message": "知识库文档上传成功"
            # "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"文档上传失败: {str(e)}")